from .tools.util_openai_sdk import ChatModelApi

__VERSION__ = "v1.0.0"

class TextTranslatorNode:
    CATEGORY = "AmorForge"
    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("translated_text",)
    OUTPUT_NODE = False

    FUNCTION = "translate"
    NODE_CLASS_MAPPINGS = {
        "TextTranslatorNode": "Text Translator Node",
    }

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "text": ("STRING", {
                    "multiline": True,
                    "default": "Hello, world!",
                    "dynamicPrompts": True
                }),
                "target_language": (["English", "Chinese", "Japanese", "Korean", "French", "German", "Spanish"], {
                    "default": "English"
                }),
            },
        }

    def translate(self, text, target_language="English"):
        """Translate text to the target language using OpenAI API"""
        try:
            # Check if text is empty
            if not text or not text.strip():
                return ("",)

            # Initialize the ChatModelApi
            translator = ChatModelApi()

            # Connect to OpenAI API
            # The API key should be set in environment variables
            try:
                translator.connect()
            except Exception as e:
                error_msg = f"Failed to connect to translation service. Please check your API key configuration: {str(e)}"
                print(f"[TextTranslatorNode] Connection Error: {error_msg}")
                return (error_msg,)

            prompt = f"You are a professional translator. Translate the following text to {target_language}. Return ONLY the translated text without any explanations, prefixes, or additional content:\n\n{text}"

            # Perform translation
            translated_text = translator.generate(prompt)

            # Clean up the response if needed
            if translated_text:
                translated_text = translated_text.strip()

            print(f"[TextTranslatorNode] Translation completed: {translated_text}")
            return (translated_text,)

        except Exception as e:
            # Return error message as translation if something goes wrong
            error_msg = f"Translation failed: {str(e)}"
            print(f"[TextTranslatorNode] Error: {error_msg}")
            return (error_msg,)
